# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/21 19:01 
@Author  : ZhangShenao 
@File    : 5.继承BaseTool创建自定义工具.py 
@Desc    : 继承BaseTool父类,创建自定义工具

通过继承BaseTool父类,并实现_run()方法,可以创建自定义工具
如果想创建自定义工具,并且目前没有任何已经实现的代码,则可以采用这种方式
"""
from typing import Any, Type

from langchain_core.tools import BaseTool
from pydantic import BaseModel

from calculator_input import CalculatorInput


class MultiplyTool(BaseTool):
    """自定义乘法工具"""

    name = "multiply_tool"  # 设置工具名称
    description = "乘法工具,输入两个整数,返回乘积"  # 设置工具描述信息
    args_schema: Type[BaseModel] = CalculatorInput  # 指定工具参数格式
    return_direct = True  # 设置工具为直接返回

    def _run(self, *args: Any, **kwargs: Any) -> Any:
        x = kwargs.get("x")
        y = kwargs.get("y")
        return x * y


# 创建自定义工具
multiply_tool = MultiplyTool()

# 获取工具的属性信息
print(f"工具名称: {multiply_tool.name}")
print(f"工具描述: {multiply_tool.description}")
print(f"工具参数定义: {multiply_tool.args}")
print(f"工具是否直接返回: {multiply_tool.return_direct}")

# 调用工具,获取结果
result = multiply_tool.invoke({"x": 3, "y": 4})
print(f"工具调用结果: {result}")
